#pragma once

#include <vector>
#include <utility>

#include "../utils/misc.h"
#include "whimed.h"

namespace mdtx::online::algo
{
    /**
     * @brief Medcouple algorithm
     *
     * @details Johnson & Mizoguchi Kth pair algorithm: https://en.wikipedia.org/wiki/Medcouple
     *
     * @tparam T
     */
    template <typename T>
    class medcouple
    {
    public:
        using tvect = std::vector<T>;
        using ivect = std::vector<ptrdiff_t>;
        using tpair = std::pair<T, ptrdiff_t>;
        using pvect = std::vector<tpair>;

    protected:
        T eps1, eps2;
        ptrdiff_t n, nplus{}, nminus{};

        tvect zplus, zminus;
        ivect left, right, p, q;
        pvect pwork;
        tvect work;

        T h_kernel(ptrdiff_t i, ptrdiff_t j) const
        {
            T a = zplus[i], b = zminus[j];
            if (a == b)
            {
                return sign(nplus - 1 - i - j);
            }
            else
            {
                return (a + b) / (a - b);
            }
        }

    public:
        medcouple(ptrdiff_t n, T eps1, T eps2)
            : eps1(eps1), eps2(eps2),
              n(n),
              zplus(n), zminus(n),
              left(n), right(n), p(n), q(n),
              pwork(n),
              work(n)
        {
        }

        template <typename iter>
        auto run(iter it, T med)
        {
            ptrdiff_t i, j, k;
            ptrdiff_t nl, nr, mcidx;

            auto get = [it](ptrdiff_t idx)
            {
                return *(it + idx);
            };

            if (std::abs(get(0) - med) < eps1 * std::abs(med))
            {
                return T(-1);
            }
            if (std::abs(get(n - 1) - med) < eps1 * std::abs(med))
            {
                return T(1);
            }

            T ztmp;
            zplus.clear();
            zminus.clear();
            for (i = 0; i < n; ++i)
            {
                ztmp = get(i) - med;
                if (std::abs(ztmp) <= eps2 * 2)
                {
                    zplus.push_back(0);
                    zminus.push_back(0);
                }
                else if (ztmp > 0)
                {
                    zplus.push_back(ztmp);
                }
                else
                {
                    zminus.push_back(ztmp);
                }
            }

            nplus = zplus.size();
            nminus = zminus.size();
            std::fill(left.begin(), left.begin() + nplus, 0);
            std::fill(right.begin(), right.begin() + nplus, nminus - 1);
            nl = 0;
            nr = nminus * nplus;
            mcidx = nr / 2;

            while (nr - nl > nplus)
            {
                j = 0;
                for (i = 0; i < nplus; ++i)
                {
                    if (left[i] <= right[i])
                    {
                        pwork[j].first = h_kernel(i, (left[i] + right[i]) / 2);
                        pwork[j].second = right[i] - left[i] + 1;
                        ++j;
                    }
                }
                T rowm = whimed::run(pwork.begin(), j);
                T rowm_eps = eps1 * std::abs(rowm);

                // P
                j = 0;
                for (i = nplus - 1; i >= 0; --i)
                {
                    while (j < nminus && h_kernel(i, j) - rowm > rowm_eps)
                    {
                        ++j;
                    }
                    p[i] = j - 1;
                }
                // Q
                j = nminus - 1;
                for (i = 0; i < nplus; ++i)
                {
                    while (j >= 0 && rowm - h_kernel(i, j) > rowm_eps)
                    {
                        --j;
                    }
                    q[i] = j + 1;
                }

                auto sump = std::accumulate(p.begin(), p.begin() + nplus, ptrdiff_t{}) + nplus;
                auto sumq = std::accumulate(q.begin(), q.begin() + nplus, ptrdiff_t{});

                if (mcidx <= sump - 1)
                {
                    std::copy(p.begin(), p.begin() + nplus, right.begin());
                    nr = sump;
                }
                else if (mcidx > sumq - 1)
                {
                    std::copy(q.begin(), q.begin() + nplus, left.begin());
                    nl = sumq;
                }
                else
                {
                    return rowm;
                }
            }

            k = 0;
            for (i = 0; i < nplus; ++i)
            {
                for (j = left[i]; j <= right[i]; ++j)
                {
                    work[k] = h_kernel(i, j);
                    ++k;
                }
            }

            mcidx -= nl;
            std::nth_element(work.begin(), work.begin() + mcidx, work.begin() + k, std::greater<T>{});
            return work[mcidx];
        }
    };
}
